import json
import os

from dora import Node, DoraStatus
import pyarrow as pa
from mofa.kernel.utils.util import load_agent_config, create_agent_output, load_node_result
from mofa.run.run_agent import run_dspy_or_crewai_agent
from mofa.utils.files.dir import get_relative_path
from mofa.utils.log.agent import record_agent_result_log


# def extract_relevant_context(context_rag, task_keywords):
#     """
#     提取与任务关键字相关的上下文。
#     根据传入的任务关键字，从context_rag中提取最相关的上下文信息。
#     """
#     relevant_contexts = []
#
#     for context in context_rag:
#         if isinstance(context, dict):
#             for key, value in context.items():
#                 # 如果任务关键字出现在context key中，提取该段落
#                 if any(kw in key.lower() for kw in task_keywords):
#                     relevant_contexts.append(" ".join(value) if isinstance(value, list) else str(value))
#         elif isinstance(context, str):
#             if any(kw in context.lower() for kw in task_keywords):
#                 relevant_contexts.append(context)
#
#     merged_context = " ".join(relevant_contexts) if relevant_contexts else "未找到相关上下文。"
#
#     print('merged_context:', merged_context)
#     return merged_context


class Operator:
    def __init__(self):
        self.task = None
        self.context_rag = None
        self.web_search = None

    def format_content_rag(self, context_rag):
        """
        将 context_rag 格式化为一个长字符串，方便大模型处理。
        """
        formatted_context = ""
        if isinstance(context_rag, list):
            # 将列表中的内容拼接成文本块
            for entry in context_rag:
                for key, value in entry.items():
                    formatted_context += f"\nTopic: {key}\nDetails: {' '.join(value)}\n"
        return formatted_context

    def on_event(self, dora_event, send_output) -> DoraStatus:
        if dora_event["type"] == "INPUT":
            if dora_event['id'] == 'task':
                self.task = dora_event["value"][0].as_py()
            if dora_event['id'] == 'context_rag':
                self.context_rag = load_node_result(dora_event["value"][0].as_py())
            if dora_event['id'] == 'web_search':
                self.web_search = dora_event["value"][0].as_py()

            if self.context_rag is not None and self.task is not None:
                # 将 task 和 context_rag 结合起来，生成完整输入
                # self.context_rag = self.format_content_rag(self.context_rag)
                # full_input = f"Task: {self.task}\n\nContext:\n{formatted_context_rag}"

                yaml_file_path = get_relative_path(
                    current_file=__file__,
                    sibling_directory_name='configs',
                    target_file_name='reasoner_agent.yml'
                )

                # 加载配置文件
                inputs = load_agent_config(yaml_file_path)

                inputs["task"] = self.task
                inputs['input_fields'] = {"content_rag":self.context_rag,"web_search":self.web_search}

                print(f"Agent Config: {inputs}")

                # print(f"Full Input: {full_input}")

                # 调用模型执行
                agent_result = run_dspy_or_crewai_agent(agent_config=inputs)


                log={"task":self.task,"response":agent_result}


                # 记录日志并返回结果
                record_agent_result_log(agent_config=inputs, agent_result=log)

                send_output("reasoner_response", pa.array([create_agent_output(
                    step_name='reasoner_response',
                    output_data=agent_result,
                    dataflow_status=os.getenv('IS_DATAFLOW_END', True)
                )]), dora_event['metadata'])

                print('reasoner_response:', agent_result)

        return DoraStatus.CONTINUE